import json
import re
from typing import Dict, List

from openai import OpenAI

from core.config import OPENAI_CONFIG

config = OPENAI_CONFIG[-1]
model = config.get("model")
client = OpenAI(
    api_key=config.get("api_key"),
    base_url=config.get("base_url"),
)


def generate(messages: List) -> str:
    """
    Generate a response from OpenAI API.
    """
    completion = client.chat.completions.create(model=model, messages=messages)
    return completion.choices[0].message.content


def stream_generate(
    messages: List,
    sentence_split: bool = False,
):
    """
    Stream chat response from OpenAI API.
    """

    def extract_complete_sentences(buffer: str):
        # 使用正则表达式提取完整句子（以句号、问号或感叹号结尾）
        parts = re.split(r'(?<=[。？！.?!])(?=[”"]?)', buffer)
        if len(parts) <= 1:
            return [], buffer
        return parts[:-1], parts[-1]

    completion = client.chat.completions.create(
        model=model, messages=messages, stream=True
    )
    string_buf = ""
    for chunk in completion:
        if not getattr(chunk, "choices", []):
            continue

        delta = chunk.choices[0].delta
        content = getattr(delta, "content", "")
        if not content:
            continue

        if not sentence_split:
            yield content
        else:
            string_buf += content
            sentences, string_buf = extract_complete_sentences(string_buf)
            if sentences:
                yield "".join(sentences)


def json_generate(messages: List) -> Dict:
    """
    Generate a response from OpenAI API with JSON output.
    """
    completion = client.chat.completions.create(
        model=model, messages=messages, response_format={"type": "json_object"}
    )
    if not getattr(completion, "choices", []):
        return {}

    message = completion.choices[0].message
    content = getattr(message, "content", "")
    try:
        return json.loads(content)
    except json.JSONDecodeError:
        print(f"Failed to decode JSON: {content}")
        return {}


def reasoning_generate(messages: List, has_reasoning: bool = False):
    completion = client.chat.completions.create(
        model=model,
        messages=messages,
        extra_body={"enable_thinking": True},
        stream=True,
    )

    for chunk in completion:
        if not getattr(chunk, "choices", []):
            continue

        delta = chunk.choices[0].delta
        reasoning = getattr(delta, "reasoning_content", "")
        content = getattr(delta, "content", "")
        if has_reasoning and reasoning:
            yield {"type": "reasoning", "content": reasoning}
        elif content:
            yield {"type": "content", "content": content}


if __name__ == "__main__":
    test_messages = [{"role": "user", "content": "Hello, how are you?"}]
    for response_chunk in reasoning_generate(test_messages, has_reasoning=True):
        print(response_chunk)
